#!/usr/bin/env python

import os, sys
from astropy.io import fits
import numpy as np
import matplotlib.pylab as plt

def main():
    if len(sys.argv) != 2:
        print('usage: {} darke.fits'.format(sys.argv[0]))
    darke = fits.getdata(sys.argv[1])
    # 先标记出所有热像元（孤立的+成列的）
    hot_index = darke > 5
    
    hot_mask = np.zeros(darke.shape, dtype=np.uint8)
    hot_mask[hot_index] = 1
    hdu = fits.PrimaryHDU(hot_mask)
    hdu.writeto('hot_mask.fits', overwrite=True)
    
    # 将OS1-8的index提取出来
    os18 = hot_index[0:4616,:]
    # 将OS9-16的index提取出来
    os916= hot_index[4616:,:]
    # 将OS1-8的index按列求和
    sum_os18 = np.sum(os18,axis=0)
    # 将OS9-16的index按列求和
    sum_os916= np.sum(os916,axis=0)

    plt.plot(sum_os18, label='OS1-8')
    plt.plot(sum_os916, label='OS9-16')
    plt.legend()
    plt.show()

    # 进一步判断那些列中热像元数目不超过400
    index_18 = sum_os18 < 400
    index_916= sum_os916 < 400

    # 将不超过100个热像元的列中的热像元数目相加
    nhots = np.sum(sum_os18[index_18]) + np.sum(sum_os916[index_916])
    print('nhots: {}'.format(nhots))
    print('hot cols: {}'.format( np.sum(sum_os18>=400) + np.sum(sum_os916>=400)))


if __name__ == '__main__':
    main()


